package com.hzw.saas.web.shell.interceptor;

import com.hzw.saas.common.security.token.SaasTokenService;
import jodd.util.StringUtil;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.util.Map;

/**
 * <p>
 *  websocket握手（handshake）接口
 * </p>
 *
 * @author sonam
 * @since 2021/1/12 17:34
 */
@Slf4j
@AllArgsConstructor
public class HttpSessionIdHandshakeInterceptor extends HttpSessionHandshakeInterceptor {

    private final SaasTokenService userTokenServices;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request,
                                   ServerHttpResponse response, WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) throws Exception {
        //解决The extension [x-webkit-deflate-frame] is not supported问题
        if(request.getHeaders().containsKey("Sec-WebSocket-Extensions")) {
            request.getHeaders().set("Sec-WebSocket-Extensions", "permessage-deflate");
        }

        if (request instanceof ServletServerHttpRequest) {
            // 拿到token自动登录
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            String token = servletRequest.getServletRequest().getParameter("token");
            if(StringUtil.isBlank(token)) {
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }

            try {
                OAuth2Authentication auth2Authentication = userTokenServices.loadAuthentication(token.replace("bearer", ""));
                SecurityContextHolder.getContext().setAuthentication(auth2Authentication);
            } catch (InvalidTokenException e) {
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }

        }
        return super.beforeHandshake(request, response, wsHandler, attributes);
    }


    @Override
    public void afterHandshake(ServerHttpRequest request,
                               ServerHttpResponse response, WebSocketHandler wsHandler,
                               Exception ex) {
        // TODO 记录在线人数？
        super.afterHandshake(request, response, wsHandler, ex);
    }


}
